package com.fc.connector.source;


import org.apache.flink.api.common.eventtime.Watermark;
import org.apache.flink.api.connector.source.*;
import org.apache.flink.core.io.InputStatus;
import org.apache.flink.table.data.GenericRowData;
import org.apache.flink.table.data.RowData;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.CompletableFuture;

public class MySourceReader implements SourceReader<RowData, MySplit> {
    // 配置参数
    private static final int NUM_RECORDS_PER_SPLIT = 5; // 每个分片模拟生成的数据量
    private static final long WATERMARK_INTERVAL_MS = 5000; // 水位线间隔

    // 运行时状态
    private final SourceReaderContext context;
    private final Queue<MySplit> assignedSplits = new LinkedList<>();
    private final Map<String, Integer> splitProgress = new HashMap<>(); // 分片消费进度
    private long nextWatermarkTime = 0L;

    // 当前处理状态
    private MySplit currentSplit;
    private int currentRecordCount = 0;

    public MySourceReader(SourceReaderContext context) {
        this.context = context;
    }

    //-------------------------- 核心方法实现 --------------------------

    /**
     * 开启SourceReader，使得SourceReader开始从数据源中读取数据。
     */
    @Override
    public void start() {
        // 初始化时向 SplitEnumerator 请求第一个分片（可选）
        context.sendSplitRequest();
    }

    /**
     * 将下一个可用的数据记录放入SourceOutput中。
     * 接口实现必须保证这个方法不会被阻塞。
     * 尽管接口实现可以提交多条记录到SourceOutput，但是不推荐这样做。相反，反而推荐提交一条记录到SourceOutput，然后返回InputStatus.MORE_AVAILABLE结果去让调用者线程知道还有更多的数据记录。
     * 参数output，接收SourceReader读取到的记录。
     * 返回参数InputStatus，方法调用后，SourceReader的InputStatus，是还有更多的记录，还是没有记录了。
     * @param output
     * @return
     * @throws Exception
     */
    @Override
    public InputStatus pollNext(ReaderOutput<RowData> output) throws Exception {
        // 1. 检查并处理新分配的分片
        if (currentSplit == null && !assignedSplits.isEmpty()) {
            currentSplit = assignedSplits.poll();
            currentRecordCount = splitProgress.getOrDefault(currentSplit.splitId(), 0);
        }

        // 2. 无可用数据时返回
        if (currentSplit == null) {
            return InputStatus.MORE_AVAILABLE;
        }

        // 3. 生成模拟数据
        if (currentRecordCount < NUM_RECORDS_PER_SPLIT) {
            // 生成记录（示例数据：分片ID的哈希值 + 递增序号）
            GenericRowData rowData = new GenericRowData(3); // 3个字段
            rowData.setField(0, currentSplit.splitId()+"_"+currentRecordCount); // 分片ID
            rowData.setField(1, currentRecordCount); // 记录序号
            rowData.setField(2, System.currentTimeMillis()); // 时间戳

            output.collect(rowData);

            currentRecordCount++;
            splitProgress.put(currentSplit.splitId(), currentRecordCount);

            // 周期发送水位线
            emitWatermarkIfNeeded(output);
            return InputStatus.MORE_AVAILABLE;
        }

        // 4. 当前分片处理完成
        finishCurrentSplit();
        return InputStatus.MORE_AVAILABLE;
    }

    //-------------------------- 分片管理 --------------------------

    /**
     * 添加一个分片列表给到数据源读取器,读取方法调用时机是在枚举器分配分片的时候调用（ 如SplitEnumeratorContext. assignSplit(SourceSplit, int) 或者 SplitEnumeratorContext. assignSplits(SplitsAssignment) ）。
     * @param splits 由枚举器分配的分片列表。
     */
    @Override
    public void addSplits(List<MySplit> splits) {
        assignedSplits.addAll(splits);
        splits.forEach(split ->
            splitProgress.putIfAbsent(split.splitId(), 0)
        );
    }


    /**
     * 这个方法当数据源读取器没有获取到更多的分片时被调用。
     * 具体触发的时机是在枚举器调用了 SplitEnumeratorContext.signalNoMoreSplits(int）方法通知数据源读取器 subtask无更多分片时调用。
     */
    @Override
    public void notifyNoMoreSplits() {
        // 可在此处理无更多分片的情况
        System.out.println("可在此处理无更多分片的情况");
    }

    //-------------------------- 状态管理 --------------------------

    /**
     * 通知checkpoint完成，这个不是SourceReader的原生接口，而是通过继承CheckpointListener接口。
     * 参数checkpointId，表示某次checkpoint的标识ID。
     * @param checkpointId
     * @return
     */
    @Override
    public List<MySplit> snapshotState(long checkpointId) {
        // 返回所有未完成的分片（包含进度）
        List<MySplit> state = new ArrayList<>();
        if (currentSplit != null) {
            state.add(currentSplit); // 当前处理中的分片
        }
        state.addAll(assignedSplits); // 待处理分片
        return state;
    }


    /**
     * 返回一个feature，feature完成，表示可以从源读取器中读取数据。
     * 如果数据源读取器中有可用的数据，那么调用此方法，需要返回一个完成的feature，否则数据源读取器将一直停滞。如果数据源读取器中无可用的数据，那么需要返回一个未完成的feature，待数据源读取器中有可用的数据后，再把feature置为已完成，数据源读取器中的数据可被读取。
     * @return
     */
    @Override
    public CompletableFuture<Void> isAvailable() {
        return null;
    }

    /**
     * 通知checkpoint完成，这个不是SourceReader的原生接口，而是通过继承CheckpointListener接口。
     * 参数checkpointId，表示某次checkpoint的标识ID。
     * @param checkpointId
     */
    @Override
    public void notifyCheckpointComplete(long checkpointId) {
        // 可选：确认检查点完成后的处理
    }

    //-------------------------- 资源管理 --------------------------
    @Override
    public void close() throws IOException {
        assignedSplits.clear();
        splitProgress.clear();
        currentSplit = null;
    }

    //-------------------------- 私有方法 --------------------------
//    private InputStatus maybeRequestMoreSplits() {
//        if (context.registeredReaders().get(context.getIndexOfSubtask()) != null) {
//            context.sendSplitRequest(); // 请求更多分片
//            return InputStatus.NOTHING_AVAILABLE;
//        }
//        return InputStatus.END_OF_INPUT;
//    }

    private void finishCurrentSplit() {
        splitProgress.remove(currentSplit.splitId());
        currentSplit = null;
        currentRecordCount = 0;
    }

    private void emitWatermarkIfNeeded(ReaderOutput<RowData> output) {
        final long now = System.currentTimeMillis();
        if (now > nextWatermarkTime) {
            output.emitWatermark(new Watermark(now));
            nextWatermarkTime = now + WATERMARK_INTERVAL_MS;
        }
    }
}